import pandas as pd
import numpy as np


def extract_dt(df):
    df['time'] = pd.to_datetime(df['time'], format='%m%d %H:%M:%S')
    # df['month'] = df['time'].dt.month
    # df['day'] = df['time'].dt.day
    df['date'] = df['time'].dt.date
    df['hour'] = df['time'].dt.hour
    # df = df.drop_duplicates(['ship','month'])
    df['weekday'] = df['time'].dt.weekday
    return df

train = pd.read_hdf('./input/train.h5')
test = pd.read_hdf('./input/test.h5')
train = extract_dt(train)
test = extract_dt(test)
ship_type = train.drop_duplicates(['ship', 'type'])[['ship', 'type']]
type_map = {'围网': 0, '拖网': 1, '刺网': 2}
type_map_rev = {v: k for k, v in type_map.items()}
ship_type['type_label'] = ship_type['type'].map(type_map)


def diff_op(df, cols_sp):
    def diff(a, cols_sp):
        col = a.name
        if col not in cols_sp:
            arr = a.tolist()
            b = pd.Series(arr[1:] + [arr[-1]]) - a
        else:
            arr = a.tolist()
            b = pd.Series(arr[1:] + [arr[-1]]) - a
        return b
    df1 = df.sort_values(by='time').reset_index().drop('index', axis=1)
    df2 = df1.apply(diff, axis=0, **{'cols_sp': cols_sp})
    return df2


# train['rnk'] = range(train.shape[0])
# cols = ['x', 'y', 'v', 'd', 'time', 'rnk', 'ship']
# feas = ['x', 'y', 'v', 'd', 'time', 'rnk']
cols = ['x', 'y', 'v', 'd', 'time', 'ship']
feas = ['x', 'y', 'v', 'd', 'time']
cols_sp = ['time']
trn = train[cols].groupby('ship')[feas].apply(diff_op, cols_sp).\
    reset_index().drop('level_1', axis=1)

trn['time'] = trn['time'].dt.seconds
trn['xxy'] = abs(trn['x'])*abs(trn['y'])
trn['x_y'] = pd.Series((trn['x']**2 + trn['y']**2)**0.5)
trn['y/x'] = trn['y']/trn['x']
trn['x/y'] = trn['x']/trn['y']
trn['xt'] = trn['x']/trn['time']
trn['yt'] = trn['y']/trn['time']
trn['vt'] = trn['v']/trn['time']
trn['dt'] = trn['d']/trn['time']
trn['xxyt'] = trn['xxy']/trn['time']
trn['x_yt'] = trn['x_y']/trn['time']
trn.fillna(0, inplace=True)

feas2 = ['xxy', 'x_y', 'y/x', 'x/y', 'xt', 'yt', 'vt', 'dt', 'xxyt', 'x_yt']
feas += feas2
trn1 = trn.drop_duplicates('ship')
desc_cols = ['count', 'mean', 'std', 'min', '25%', '50%', '75%', 'max']
for fea in feas:
    trn[fea][abs(trn[fea]) > 10*10] = 0
    tmp = trn.groupby('ship')[fea].describe().reset_index()
    tmp.columns = ['ship'] + [fea + '_diff_' + a for a in desc_cols if a not in ['ship']]
    if fea in ['xxyt', 'x_yt', 'x_y', 'xxy']:
        tmp[fea + '_sum'] = trn.groupby('ship')[fea].sum()
    trn1 = pd.merge(trn1, tmp, on='ship', how='left', suffixes=('', ''))

trn1['type'] = trn1['ship'].map(ship_type.set_index('ship')['type_label'])
features = [a for a in trn1.columns if a not in ['ship', 'type']] + ['ship', 'type']
trn2 = trn1[features]
trn2.fillna(0, inplace=True)
trn2.replace({np.inf: 10**10, -np.inf: -10**10}, inplace=True)
trn2.values[np.where(abs(trn2.values) > 10**10)]
trn2.to_csv('data/X_desc_diff.csv', header=True, index=False)


outfeas = [a for a in trn2.columns if a not in ['ship', 'type']]
import matplotlib.pyplot as plt
for fea in outfeas:
    plt.figure()
    plt.plot(trn2['type'], trn2[fea], '*')
    # plt.show()
    fea1 = fea.replace('/', '__')
    plt.savefig('feas/'+fea1+'.png')

fea2 = []
N0 = []
for fea in outfeas:
    num_u = trn2[fea].nunique()
    num_0 = sum(trn2[fea] == 0)
    if num_u < 5 or num_0 > 6000:
        print(fea, num_u, num_0)
    else:
        fea2 += [fea]

trn3 = trn2[fea2+['ship', 'type']]
trn3.to_csv('data/X_desc_diff_0nu.csv', header=True, index=False)

